
import torch

# batch size 维度设置为 1， 方便 torch 进行广播
# Output : (1, query_len, key_len)

# Usage : 根据输入的 query, key 生成对应的 上三角 mask tensor (1, query, key)
#         用来 掩盖掉 decoder 输入 中 的 对应答案，比如
#         tensor([[[ True, False, False, False, False],
#                  [ True,  True, False, False, False],
#                  [ True,  True,  True, False, False],
#                  [ True,  True,  True,  True, False],
#                  [ True,  True,  True,  True,  True]]])
#         mask.shape = (batch_size, query_len, key_len) query 行， key 列
#         其中 mask[i, j] 表示：第 i 个 query 能否看到第 j 个 key
# @param (query) : query 的大小
# @param (key) : key 的大小
def subsequentMask(query : int, key : int):
    # torch.triu(..., diagonal=1) 取张量的 上三角部分，不包含主对角线。
    # diagonal=1 表示从第 1 条副对角线开始（即严格的上三角）
    # .bool() 表示 把 float 类型的 0/1 张量转换成布尔类型 False/True
    # torch.logical_not(...) 用于取反
    return  torch.logical_not(
                torch.triu(torch.ones(1, query, key), diagonal = 1).bool()
            )


if __name__ == "__main__":
    print(subsequentMask(5, 5))
# tensor([[[ True, False, False, False, False],
#          [ True,  True, False, False, False],
#          [ True,  True,  True, False, False],
#          [ True,  True,  True,  True, False],
#          [ True,  True,  True,  True,  True]]])

    print(subsequentMask(6, 4))
# tensor([[[ True, False, False, False],
#          [ True,  True, False, False],
#          [ True,  True,  True, False],
#          [ True,  True,  True,  True],
#          [ True,  True,  True,  True],
#          [ True,  True,  True,  True]]])
